# -*- coding: utf-8 -*-
# !/usr/bin/env python
"""
-------------------------------------------------
   File Name：     test
   Description :   
   Author :       lth
   date：          2022/6/10
-------------------------------------------------
   Change Activity:
                   2022/6/10 10:36: create this script
-------------------------------------------------
"""
__author__ = 'lth'

import copy

import torch
from PIL import Image

from datalist import val_transform
from model import AnimeGenerator
from pretrain import denormalize


class Inference:
    def __init__(self):
        self.generator = AnimeGenerator().cuda()
        self.generator.eval()
        model_state_dict = torch.load("weights/generate_best.pth")["model_state_dict"]
        new_state_dict = {}
        for k, v in model_state_dict.items():
            new_state_dict[k.replace("module.", "")] = v

        self.generator.load_state_dict(new_state_dict)

    @torch.no_grad()
    def inference(self, image):

        height = image.height
        width = image.width

        height = height // 32 * 32
        width = width // 32 * 32

        image = image.resize((width, height))

        image = val_transform(image).unsqueeze(0).cuda()

        generator_img = self.generator(image)
        generator_img = generator_img * 0.9 + image * 0.1

        self.get_image(generator_img, image)

    @staticmethod
    def get_image(image, photo):
        output = (denormalize(image.permute((0, 2, 3, 1)).
                              detach().
                              to('cpu').
                              numpy()) * 255).astype('uint8')
        output = output[0]

        photo = (denormalize(photo.permute((0, 2, 3, 1)).
                             detach().
                             to('cpu').
                             numpy()) * 255).astype('uint8')

        photo = photo[0]

        width = photo.shape[1]
        height = photo.shape[0]

        output = Image.fromarray(output).convert('RGB')
        photo = Image.fromarray(photo).convert("RGB")

        target = Image.new('RGB', (width + width, height), (255, 255, 255))

        target.paste(output, (0, 0, width, height))
        target.paste(photo, (width, 0, 2 * width, height))

        target.show()
        target.save("save.jpg")

    @torch.no_grad()
    def inference_from_video(self, video_path):

        import cv2
        import numpy as np
        import time
        capture = cv2.VideoCapture(video_path)

        fps = 0
        while True:
            t1 = time.time()
            ref, frame = capture.read()
            temp = copy.deepcopy(frame)
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

            frame = val_transform(frame).unsqueeze(0).cuda()
            generator_img = self.generator(frame)
            generator_img=generator_img*0.9+frame*0.1

            output = (denormalize(generator_img.permute((0, 2, 3, 1)).
                                  detach().
                                  to('cpu').
                                  numpy()) * 255).astype('uint8')
            frame = output[0]
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame = np.concatenate([frame, temp])

            fps = (fps + (1. / (time.time() - t1))) / 2
            print("fps= %.2f" % (fps))
            frame = cv2.putText(frame, "fps= %.2f" % (fps), (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

            cv2.imshow("video", frame)
            c = cv2.waitKey(1) & 0xff

            if c == 27:
                capture.release()
                break
        capture.release()
        cv2.destroyAllWindows()


if __name__ == "__main__":
    model = Inference()
    img = Image.open("1.png").convert("RGB")
    # model.inference(img)
    model.inference_from_video(0)
